{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 量化分区"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "构建 PyTorch 前端模型："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch import nn\n",
    "import torch\n",
    "\n",
    "\n",
    "class Model(nn.Module):\n",
    "    def __init__(self, *args, **kwargs) -> None:\n",
    "        super().__init__(*args, **kwargs)\n",
    "        self.conv = nn.Conv2d(3, 16, 3, 1, 1, bias=False)\n",
    "        self.conv2 = nn.Conv2d(16, 16, 3, 1, 1, bias=True)\n",
    "        self.relu = nn.ReLU()\n",
    "        self.pool = nn.MaxPool2d(kernel_size=2)\n",
    "        self.dense1 = nn.Linear(64, 32, bias=False)\n",
    "        self.dense2 = nn.Linear(32, 16)\n",
    "        self.dense3 = nn.Linear(16, 8)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.conv(x)\n",
    "        x = self.relu(x)\n",
    "        x = self.conv2(x)\n",
    "        x = self.relu(x)\n",
    "        x = self.pool(x)\n",
    "        x = x.flatten(1)\n",
    "        x = self.dense1(x)\n",
    "        x = self.relu(x)\n",
    "        x = self.dense2(x)\n",
    "        x = self.relu(x)\n",
    "        x = self.dense3(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "翻译 PyTorch 前端模型为 Relay 模型："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import set_env"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fn (%data: Tensor[(1, 3, 4, 4), float32] /* ty=Tensor[(1, 3, 4, 4), float32] span=aten___convolution_0_data:0:0 */) -> Tensor[(1, 8), float32] {\n",
      "  %0 = nn.conv2d(%data, meta[relay.Constant][0] /* ty=Tensor[(16, 3, 3, 3), float32] */, padding=[1, 1, 1, 1], channels=16, kernel_size=[3, 3]) /* ty=Tensor[(1, 16, 4, 4), float32] span=aten___convolution_0:0:0 */;\n",
      "  %1 = nn.relu(%0) /* ty=Tensor[(1, 16, 4, 4), float32] span=aten__relu_0:0:0 */;\n",
      "  %2 = nn.conv2d(%1, meta[relay.Constant][1] /* ty=Tensor[(16, 16, 3, 3), float32] */, padding=[1, 1, 1, 1], channels=16, kernel_size=[3, 3]) /* ty=Tensor[(1, 16, 4, 4), float32] span=aten___convolution_1:0:0 */;\n",
      "  %3 = add(%2, meta[relay.Constant][2] /* ty=Tensor[(16, 1, 1), float32] */) /* ty=Tensor[(1, 16, 4, 4), float32] */;\n",
      "  %4 = nn.relu(%3) /* ty=Tensor[(1, 16, 4, 4), float32] span=aten__relu_1:0:0 */;\n",
      "  %5 = nn.max_pool2d(%4, pool_size=[2, 2], strides=[2, 2], padding=[0, 0, 0, 0]) /* ty=Tensor[(1, 16, 2, 2), float32] span=aten__max_pool2d_0:0:0 */;\n",
      "  %6 = reshape(%5, newshape=[0, -1, 1, 1]) /* ty=Tensor[(1, 64, 1, 1), float32] span=aten__flatten_0:0:0 */;\n",
      "  %7 = squeeze(%6, axis=[2, 3]) /* ty=Tensor[(1, 64), float32] span=aten__flatten_0:0:0 */;\n",
      "  %8 = nn.dense(%7, meta[relay.Constant][3] /* ty=Tensor[(32, 64), float32] */, units=None) /* ty=Tensor[(1, 32), float32] span=aten__linear_0:0:0 */;\n",
      "  %9 = nn.relu(%8) /* ty=Tensor[(1, 32), float32] span=aten__relu_2:0:0 */;\n",
      "  %10 = nn.dense(%9, meta[relay.Constant][4] /* ty=Tensor[(16, 32), float32] */, units=None) /* ty=Tensor[(1, 16), float32] span=aten__linear_1:0:0 */;\n",
      "  %11 = add(%10, meta[relay.Constant][5] /* ty=Tensor[(16), float32] */) /* ty=Tensor[(1, 16), float32] */;\n",
      "  %12 = nn.relu(%11) /* ty=Tensor[(1, 16), float32] span=aten__relu_3:0:0 */;\n",
      "  %13 = nn.dense(%12, meta[relay.Constant][6] /* ty=Tensor[(8, 16), float32] */, units=None) /* ty=Tensor[(1, 8), float32] span=aten__linear_2:0:0 */;\n",
      "  add(%13, meta[relay.Constant][7] /* ty=Tensor[(8), float32] */) /* ty=Tensor[(1, 8), float32] */\n",
      "} /* ty=fn (Tensor[(1, 3, 4, 4), float32]) -> Tensor[(1, 8), float32] */\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import set_env\n",
    "import tvm\n",
    "from tvm import relay\n",
    "\n",
    "# 输入数据\n",
    "input_shape = (1, 3, 4, 4)\n",
    "input_dtype = \"float32\"\n",
    "data_np = np.random.rand(*input_shape).astype(input_dtype)\n",
    "with torch.no_grad():\n",
    "    pt_model = Model().eval().float()\n",
    "    traced_model = torch.jit.trace(pt_model, torch.from_numpy(data_np)).eval()\n",
    "mod, params = relay.frontend.from_pytorch(traced_model, [(\"data\", input_shape)], \n",
    "                                          use_parser_friendly_name=True)\n",
    "with tvm.transform.PassContext(opt_level=3):\n",
    "    mod = relay.quantize.prerequisite_optimize(mod, params)\n",
    "print(mod['main'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "调用分区 pass："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "fn (%data: Tensor[(1, 3, 4, 4), float32] /* ty=Tensor[(1, 3, 4, 4), float32] span=aten___convolution_0_data:0:0 */) -> Tensor[(1, 8), float32] {\n",
       "  %0 = nn.conv2d(%data, meta[relay.Constant][0] /* ty=Tensor[(16, 3, 3, 3), float32] */, padding=[1, 1, 1, 1], channels=16, kernel_size=[3, 3]) /* ty=Tensor[(1, 16, 4, 4), float32] span=aten___convolution_0:0:0 */;\n",
       "  %1 = nn.relu(%0) /* ty=Tensor[(1, 16, 4, 4), float32] span=aten__relu_0:0:0 */;\n",
       "  %2 = annotation.cast_hint(%1, dtype=\"int8\") /* ty=Tensor[(1, 16, 4, 4), float32] */;\n",
       "  %3 = annotation.stop_fusion(%2) /* ty=Tensor[(1, 16, 4, 4), float32] */;\n",
       "  %4 = nn.conv2d(%3, meta[relay.Constant][1] /* ty=Tensor[(16, 16, 3, 3), float32] */, padding=[1, 1, 1, 1], channels=16, kernel_size=[3, 3]) /* ty=Tensor[(1, 16, 4, 4), float32] span=aten___convolution_1:0:0 */;\n",
       "  %5 = add(%4, meta[relay.Constant][2] /* ty=Tensor[(16, 1, 1), float32] */) /* ty=Tensor[(1, 16, 4, 4), float32] */;\n",
       "  %6 = nn.relu(%5) /* ty=Tensor[(1, 16, 4, 4), float32] span=aten__relu_1:0:0 */;\n",
       "  %7 = nn.max_pool2d(%6, pool_size=[2, 2], strides=[2, 2], padding=[0, 0, 0, 0]) /* ty=Tensor[(1, 16, 2, 2), float32] span=aten__max_pool2d_0:0:0 */;\n",
       "  %8 = annotation.cast_hint(%7, dtype=\"int8\") /* ty=Tensor[(1, 16, 2, 2), float32] */;\n",
       "  %9 = annotation.stop_fusion(%8) /* ty=Tensor[(1, 16, 2, 2), float32] */;\n",
       "  %10 = reshape(%9, newshape=[0, -1, 1, 1]) /* ty=Tensor[(1, 64, 1, 1), float32] span=aten__flatten_0:0:0 */;\n",
       "  %11 = squeeze(%10, axis=[2, 3]) /* ty=Tensor[(1, 64), float32] span=aten__flatten_0:0:0 */;\n",
       "  %12 = nn.dense(%11, meta[relay.Constant][3] /* ty=Tensor[(32, 64), float32] */, units=None) /* ty=Tensor[(1, 32), float32] span=aten__linear_0:0:0 */;\n",
       "  %13 = nn.relu(%12) /* ty=Tensor[(1, 32), float32] span=aten__relu_2:0:0 */;\n",
       "  %14 = nn.dense(%13, meta[relay.Constant][4] /* ty=Tensor[(16, 32), float32] */, units=None) /* ty=Tensor[(1, 16), float32] span=aten__linear_1:0:0 */;\n",
       "  %15 = add(%14, meta[relay.Constant][5] /* ty=Tensor[(16), float32] */) /* ty=Tensor[(1, 16), float32] */;\n",
       "  %16 = nn.relu(%15) /* ty=Tensor[(1, 16), float32] span=aten__relu_3:0:0 */;\n",
       "  %17 = nn.dense(%16, meta[relay.Constant][6] /* ty=Tensor[(8, 16), float32] */, units=None) /* ty=Tensor[(1, 8), float32] span=aten__linear_2:0:0 */;\n",
       "  add(%17, meta[relay.Constant][7] /* ty=Tensor[(8), float32] */) /* ty=Tensor[(1, 8), float32] */\n",
       "} /* ty=fn (Tensor[(1, 3, 4, 4), float32]) -> Tensor[(1, 8), float32] */"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "relay.quantize.partition()(mod)[\"main\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "为了展示分区的效果，将 TVM 默认注册的属性函数清除："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tvm.relay.op import op as _op\n",
    "def reset_partition():\n",
    "    op_names = [\"nn.conv2d\", \"nn.relu\", \"nn.max_pool2d\", \"add\", \"multiply\", \"clip\", \"nn.global_avg_pool2d\"]\n",
    "    for op_name in op_names:\n",
    "        _op.get(op_name).reset_attr(\"FQPartitionRewrite\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "此时调用分区 pass，则有："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "fn (%data: Tensor[(1, 3, 4, 4), float32] /* ty=Tensor[(1, 3, 4, 4), float32] span=aten___convolution_0_data:0:0 */) -> Tensor[(1, 8), float32] {\n",
       "  %0 = nn.conv2d(%data, meta[relay.Constant][0] /* ty=Tensor[(16, 3, 3, 3), float32] */, padding=[1, 1, 1, 1], channels=16, kernel_size=[3, 3]) /* ty=Tensor[(1, 16, 4, 4), float32] span=aten___convolution_0:0:0 */;\n",
       "  %1 = nn.relu(%0) /* ty=Tensor[(1, 16, 4, 4), float32] span=aten__relu_0:0:0 */;\n",
       "  %2 = nn.conv2d(%1, meta[relay.Constant][1] /* ty=Tensor[(16, 16, 3, 3), float32] */, padding=[1, 1, 1, 1], channels=16, kernel_size=[3, 3]) /* ty=Tensor[(1, 16, 4, 4), float32] span=aten___convolution_1:0:0 */;\n",
       "  %3 = add(%2, meta[relay.Constant][2] /* ty=Tensor[(16, 1, 1), float32] */) /* ty=Tensor[(1, 16, 4, 4), float32] */;\n",
       "  %4 = nn.relu(%3) /* ty=Tensor[(1, 16, 4, 4), float32] span=aten__relu_1:0:0 */;\n",
       "  %5 = nn.max_pool2d(%4, pool_size=[2, 2], strides=[2, 2], padding=[0, 0, 0, 0]) /* ty=Tensor[(1, 16, 2, 2), float32] span=aten__max_pool2d_0:0:0 */;\n",
       "  %6 = reshape(%5, newshape=[0, -1, 1, 1]) /* ty=Tensor[(1, 64, 1, 1), float32] span=aten__flatten_0:0:0 */;\n",
       "  %7 = squeeze(%6, axis=[2, 3]) /* ty=Tensor[(1, 64), float32] span=aten__flatten_0:0:0 */;\n",
       "  %8 = nn.dense(%7, meta[relay.Constant][3] /* ty=Tensor[(32, 64), float32] */, units=None) /* ty=Tensor[(1, 32), float32] span=aten__linear_0:0:0 */;\n",
       "  %9 = nn.relu(%8) /* ty=Tensor[(1, 32), float32] span=aten__relu_2:0:0 */;\n",
       "  %10 = nn.dense(%9, meta[relay.Constant][4] /* ty=Tensor[(16, 32), float32] */, units=None) /* ty=Tensor[(1, 16), float32] span=aten__linear_1:0:0 */;\n",
       "  %11 = add(%10, meta[relay.Constant][5] /* ty=Tensor[(16), float32] */) /* ty=Tensor[(1, 16), float32] */;\n",
       "  %12 = nn.relu(%11) /* ty=Tensor[(1, 16), float32] span=aten__relu_3:0:0 */;\n",
       "  %13 = nn.dense(%12, meta[relay.Constant][6] /* ty=Tensor[(8, 16), float32] */, units=None) /* ty=Tensor[(1, 8), float32] span=aten__linear_2:0:0 */;\n",
       "  add(%13, meta[relay.Constant][7] /* ty=Tensor[(8), float32] */) /* ty=Tensor[(1, 8), float32] */\n",
       "} /* ty=fn (Tensor[(1, 3, 4, 4), float32]) -> Tensor[(1, 8), float32] */"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "reset_partition()\n",
    "relay.quantize.partition()(mod)[\"main\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "下面看看如何从算子层面考虑对模型进行分区。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tvm.relay.quantize._partition import (\n",
    "    partition_expr_check,\n",
    "    QPartitionExpr,\n",
    "    register_partition_function\n",
    ")\n",
    "from tvm.relay.quantize.quantize import _forward_op"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[31mSignature:\u001b[39m register_partition_function(op_name, frewrite=\u001b[38;5;28;01mNone\u001b[39;00m, level=\u001b[32m10\u001b[39m)\n",
      "\u001b[31mDocstring:\u001b[39m <no docstring>\n",
      "\u001b[31mSource:\u001b[39m   \n",
      "\u001b[38;5;28;01mdef\u001b[39;00m register_partition_function(op_name, frewrite=\u001b[38;5;28;01mNone\u001b[39;00m, level=\u001b[32m10\u001b[39m):\n",
      "    \u001b[38;5;28;01mreturn\u001b[39;00m tvm.ir.register_op_attr(op_name, \u001b[33m\"FQPartitionRewrite\"\u001b[39m, frewrite, level)\n",
      "\u001b[31mFile:\u001b[39m      /media/pc/data/board/arria10/lxw/tasks/tvm-ai/python/tvm/relay/quantize/_partition.py\n",
      "\u001b[31mType:\u001b[39m      function"
     ]
    }
   ],
   "source": [
    "register_partition_function??"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "{func}`register_partition_function` 为算子注册 `\"FQPartitionRewrite\"` 属性。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[31mSignature:\u001b[39m partition_expr_check(expr)\n",
      "\u001b[31mDocstring:\u001b[39m <no docstring>\n",
      "\u001b[31mSource:\u001b[39m   \n",
      "\u001b[38;5;28;01mdef\u001b[39;00m partition_expr_check(expr):\n",
      "    \u001b[38;5;28;01mif\u001b[39;00m isinstance(expr, QPartitionExpr):\n",
      "        \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mTrue\u001b[39;00m, expr.expr\n",
      "    \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mFalse\u001b[39;00m, expr\n",
      "\u001b[31mFile:\u001b[39m      /media/pc/data/board/arria10/lxw/tasks/tvm-ai/python/tvm/relay/quantize/_partition.py\n",
      "\u001b[31mType:\u001b[39m      function"
     ]
    }
   ],
   "source": [
    "partition_expr_check??"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "{func}`partition_expr_check` 用于检查 `expr` 是否为 `QPartitionExpr`。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[31mSignature:\u001b[39m _forward_op(ref_call, args)\n",
      "\u001b[31mSource:\u001b[39m   \n",
      "\u001b[38;5;28;01mdef\u001b[39;00m _forward_op(ref_call, args):\n",
      "    \u001b[33m\"\"\"forward the operator of ref_call with provided arguments\"\"\"\u001b[39m\n",
      "    \u001b[38;5;28;01mreturn\u001b[39;00m _expr.Call(ref_call.op, args, ref_call.attrs, ref_call.type_args, ref_call.span)\n",
      "\u001b[31mFile:\u001b[39m      /media/pc/data/board/arria10/lxw/tasks/tvm-ai/python/tvm/relay/quantize/quantize.py\n",
      "\u001b[31mType:\u001b[39m      function"
     ]
    }
   ],
   "source": [
    "_forward_op??"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "{func}`_forward_op` 用于重写回调函数。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## `nn.relu` 和 `nn.max_pool2d` 分区"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "对 `nn.relu` 进行分区，可以："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tvm.relay.op import op as _op\n",
    "_op.get(\"nn.relu\").reset_attr(\"FQPartitionRewrite\")\n",
    "\n",
    "@register_partition_function(\"nn.relu\")\n",
    "def _relu_partition(ref_call, new_args, ctx):\n",
    "    return QPartitionExpr(_forward_op(ref_call, new_args))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "fn (%data: Tensor[(1, 3, 4, 4), float32] /* ty=Tensor[(1, 3, 4, 4), float32] span=aten___convolution_0_data:0:0 */) -> Tensor[(1, 8), float32] {\n",
       "  %0 = nn.conv2d(%data, meta[relay.Constant][0] /* ty=Tensor[(16, 3, 3, 3), float32] */, padding=[1, 1, 1, 1], channels=16, kernel_size=[3, 3]) /* ty=Tensor[(1, 16, 4, 4), float32] span=aten___convolution_0:0:0 */;\n",
       "  %1 = nn.relu(%0) /* ty=Tensor[(1, 16, 4, 4), float32] span=aten__relu_0:0:0 */;\n",
       "  %2 = annotation.cast_hint(%1, dtype=\"int8\") /* ty=Tensor[(1, 16, 4, 4), float32] */;\n",
       "  %3 = annotation.stop_fusion(%2) /* ty=Tensor[(1, 16, 4, 4), float32] */;\n",
       "  %4 = nn.conv2d(%3, meta[relay.Constant][1] /* ty=Tensor[(16, 16, 3, 3), float32] */, padding=[1, 1, 1, 1], channels=16, kernel_size=[3, 3]) /* ty=Tensor[(1, 16, 4, 4), float32] span=aten___convolution_1:0:0 */;\n",
       "  %5 = add(%4, meta[relay.Constant][2] /* ty=Tensor[(16, 1, 1), float32] */) /* ty=Tensor[(1, 16, 4, 4), float32] */;\n",
       "  %6 = nn.relu(%5) /* ty=Tensor[(1, 16, 4, 4), float32] span=aten__relu_1:0:0 */;\n",
       "  %7 = annotation.cast_hint(%6, dtype=\"int8\") /* ty=Tensor[(1, 16, 4, 4), float32] */;\n",
       "  %8 = annotation.stop_fusion(%7) /* ty=Tensor[(1, 16, 4, 4), float32] */;\n",
       "  %9 = nn.max_pool2d(%8, pool_size=[2, 2], strides=[2, 2], padding=[0, 0, 0, 0]) /* ty=Tensor[(1, 16, 2, 2), float32] span=aten__max_pool2d_0:0:0 */;\n",
       "  %10 = reshape(%9, newshape=[0, -1, 1, 1]) /* ty=Tensor[(1, 64, 1, 1), float32] span=aten__flatten_0:0:0 */;\n",
       "  %11 = squeeze(%10, axis=[2, 3]) /* ty=Tensor[(1, 64), float32] span=aten__flatten_0:0:0 */;\n",
       "  %12 = nn.dense(%11, meta[relay.Constant][3] /* ty=Tensor[(32, 64), float32] */, units=None) /* ty=Tensor[(1, 32), float32] span=aten__linear_0:0:0 */;\n",
       "  %13 = nn.relu(%12) /* ty=Tensor[(1, 32), float32] span=aten__relu_2:0:0 */;\n",
       "  %14 = annotation.cast_hint(%13, dtype=\"int8\") /* ty=Tensor[(1, 32), float32] */;\n",
       "  %15 = annotation.stop_fusion(%14) /* ty=Tensor[(1, 32), float32] */;\n",
       "  %16 = nn.dense(%15, meta[relay.Constant][4] /* ty=Tensor[(16, 32), float32] */, units=None) /* ty=Tensor[(1, 16), float32] span=aten__linear_1:0:0 */;\n",
       "  %17 = add(%16, meta[relay.Constant][5] /* ty=Tensor[(16), float32] */) /* ty=Tensor[(1, 16), float32] */;\n",
       "  %18 = nn.relu(%17) /* ty=Tensor[(1, 16), float32] span=aten__relu_3:0:0 */;\n",
       "  %19 = annotation.cast_hint(%18, dtype=\"int8\") /* ty=Tensor[(1, 16), float32] */;\n",
       "  %20 = annotation.stop_fusion(%19) /* ty=Tensor[(1, 16), float32] */;\n",
       "  %21 = nn.dense(%20, meta[relay.Constant][6] /* ty=Tensor[(8, 16), float32] */, units=None) /* ty=Tensor[(1, 8), float32] span=aten__linear_2:0:0 */;\n",
       "  add(%21, meta[relay.Constant][7] /* ty=Tensor[(8), float32] */) /* ty=Tensor[(1, 8), float32] */\n",
       "} /* ty=fn (Tensor[(1, 3, 4, 4), float32]) -> Tensor[(1, 8), float32] */"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "run_mod = relay.quantize.partition()(mod)\n",
    "run_mod[\"main\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "也可以对 `nn.max_pool2d` 执行同样的操作："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tvm.relay.op import op as _op\n",
    "_op.get(\"nn.relu\").reset_attr(\"FQPartitionRewrite\")\n",
    "_op.get(\"nn.max_pool2d\").reset_attr(\"FQPartitionRewrite\")\n",
    "@register_partition_function(\"nn.max_pool2d\")\n",
    "def _max_pool2d_partition(ref_call, new_args, ctx):\n",
    "    return QPartitionExpr(_forward_op(ref_call, new_args))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "fn (%data: Tensor[(1, 3, 4, 4), float32] /* ty=Tensor[(1, 3, 4, 4), float32] span=aten___convolution_0_data:0:0 */) -> Tensor[(1, 8), float32] {\n",
       "  %0 = nn.conv2d(%data, meta[relay.Constant][0] /* ty=Tensor[(16, 3, 3, 3), float32] */, padding=[1, 1, 1, 1], channels=16, kernel_size=[3, 3]) /* ty=Tensor[(1, 16, 4, 4), float32] span=aten___convolution_0:0:0 */;\n",
       "  %1 = nn.relu(%0) /* ty=Tensor[(1, 16, 4, 4), float32] span=aten__relu_0:0:0 */;\n",
       "  %2 = nn.conv2d(%1, meta[relay.Constant][1] /* ty=Tensor[(16, 16, 3, 3), float32] */, padding=[1, 1, 1, 1], channels=16, kernel_size=[3, 3]) /* ty=Tensor[(1, 16, 4, 4), float32] span=aten___convolution_1:0:0 */;\n",
       "  %3 = add(%2, meta[relay.Constant][2] /* ty=Tensor[(16, 1, 1), float32] */) /* ty=Tensor[(1, 16, 4, 4), float32] */;\n",
       "  %4 = nn.relu(%3) /* ty=Tensor[(1, 16, 4, 4), float32] span=aten__relu_1:0:0 */;\n",
       "  %5 = nn.max_pool2d(%4, pool_size=[2, 2], strides=[2, 2], padding=[0, 0, 0, 0]) /* ty=Tensor[(1, 16, 2, 2), float32] span=aten__max_pool2d_0:0:0 */;\n",
       "  %6 = annotation.cast_hint(%5, dtype=\"int8\") /* ty=Tensor[(1, 16, 2, 2), float32] */;\n",
       "  %7 = annotation.stop_fusion(%6) /* ty=Tensor[(1, 16, 2, 2), float32] */;\n",
       "  %8 = reshape(%7, newshape=[0, -1, 1, 1]) /* ty=Tensor[(1, 64, 1, 1), float32] span=aten__flatten_0:0:0 */;\n",
       "  %9 = squeeze(%8, axis=[2, 3]) /* ty=Tensor[(1, 64), float32] span=aten__flatten_0:0:0 */;\n",
       "  %10 = nn.dense(%9, meta[relay.Constant][3] /* ty=Tensor[(32, 64), float32] */, units=None) /* ty=Tensor[(1, 32), float32] span=aten__linear_0:0:0 */;\n",
       "  %11 = nn.relu(%10) /* ty=Tensor[(1, 32), float32] span=aten__relu_2:0:0 */;\n",
       "  %12 = nn.dense(%11, meta[relay.Constant][4] /* ty=Tensor[(16, 32), float32] */, units=None) /* ty=Tensor[(1, 16), float32] span=aten__linear_1:0:0 */;\n",
       "  %13 = add(%12, meta[relay.Constant][5] /* ty=Tensor[(16), float32] */) /* ty=Tensor[(1, 16), float32] */;\n",
       "  %14 = nn.relu(%13) /* ty=Tensor[(1, 16), float32] span=aten__relu_3:0:0 */;\n",
       "  %15 = nn.dense(%14, meta[relay.Constant][6] /* ty=Tensor[(8, 16), float32] */, units=None) /* ty=Tensor[(1, 8), float32] span=aten__linear_2:0:0 */;\n",
       "  add(%15, meta[relay.Constant][7] /* ty=Tensor[(8), float32] */) /* ty=Tensor[(1, 8), float32] */\n",
       "} /* ty=fn (Tensor[(1, 3, 4, 4), float32]) -> Tensor[(1, 8), float32] */"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "run_mod = relay.quantize.partition()(mod)\n",
    "run_mod[\"main\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "{func}`_max_pool2d_partition` 和 {func}`_relu_partition` 本质上是同样的功能实现，可以将它们合并为一个函数："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<function __main__.identity_partition_function(ref_call, new_args, ctx)>"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def identity_partition_function(ref_call, new_args, ctx):\n",
    "    cond, expr = partition_expr_check(new_args[0])\n",
    "    return QPartitionExpr(_forward_op(ref_call, [expr]))\n",
    "reset_partition()\n",
    "register_partition_function(\"nn.relu\", identity_partition_function)\n",
    "register_partition_function(\"nn.max_pool2d\", identity_partition_function)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "可以看出 `nn.max_pool2d` + `nn.relu` 实现了融合："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "fn (%data: Tensor[(1, 3, 4, 4), float32] /* ty=Tensor[(1, 3, 4, 4), float32] span=aten___convolution_0_data:0:0 */) -> Tensor[(1, 8), float32] {\n",
       "  %0 = nn.conv2d(%data, meta[relay.Constant][0] /* ty=Tensor[(16, 3, 3, 3), float32] */, padding=[1, 1, 1, 1], channels=16, kernel_size=[3, 3]) /* ty=Tensor[(1, 16, 4, 4), float32] span=aten___convolution_0:0:0 */;\n",
       "  %1 = nn.relu(%0) /* ty=Tensor[(1, 16, 4, 4), float32] span=aten__relu_0:0:0 */;\n",
       "  %2 = annotation.cast_hint(%1, dtype=\"int8\") /* ty=Tensor[(1, 16, 4, 4), float32] */;\n",
       "  %3 = annotation.stop_fusion(%2) /* ty=Tensor[(1, 16, 4, 4), float32] */;\n",
       "  %4 = nn.conv2d(%3, meta[relay.Constant][1] /* ty=Tensor[(16, 16, 3, 3), float32] */, padding=[1, 1, 1, 1], channels=16, kernel_size=[3, 3]) /* ty=Tensor[(1, 16, 4, 4), float32] span=aten___convolution_1:0:0 */;\n",
       "  %5 = add(%4, meta[relay.Constant][2] /* ty=Tensor[(16, 1, 1), float32] */) /* ty=Tensor[(1, 16, 4, 4), float32] */;\n",
       "  %6 = nn.relu(%5) /* ty=Tensor[(1, 16, 4, 4), float32] span=aten__relu_1:0:0 */;\n",
       "  %7 = nn.max_pool2d(%6, pool_size=[2, 2], strides=[2, 2], padding=[0, 0, 0, 0]) /* ty=Tensor[(1, 16, 2, 2), float32] span=aten__max_pool2d_0:0:0 */;\n",
       "  %8 = annotation.cast_hint(%7, dtype=\"int8\") /* ty=Tensor[(1, 16, 2, 2), float32] */;\n",
       "  %9 = annotation.stop_fusion(%8) /* ty=Tensor[(1, 16, 2, 2), float32] */;\n",
       "  %10 = reshape(%9, newshape=[0, -1, 1, 1]) /* ty=Tensor[(1, 64, 1, 1), float32] span=aten__flatten_0:0:0 */;\n",
       "  %11 = squeeze(%10, axis=[2, 3]) /* ty=Tensor[(1, 64), float32] span=aten__flatten_0:0:0 */;\n",
       "  %12 = nn.dense(%11, meta[relay.Constant][3] /* ty=Tensor[(32, 64), float32] */, units=None) /* ty=Tensor[(1, 32), float32] span=aten__linear_0:0:0 */;\n",
       "  %13 = nn.relu(%12) /* ty=Tensor[(1, 32), float32] span=aten__relu_2:0:0 */;\n",
       "  %14 = annotation.cast_hint(%13, dtype=\"int8\") /* ty=Tensor[(1, 32), float32] */;\n",
       "  %15 = annotation.stop_fusion(%14) /* ty=Tensor[(1, 32), float32] */;\n",
       "  %16 = nn.dense(%15, meta[relay.Constant][4] /* ty=Tensor[(16, 32), float32] */, units=None) /* ty=Tensor[(1, 16), float32] span=aten__linear_1:0:0 */;\n",
       "  %17 = add(%16, meta[relay.Constant][5] /* ty=Tensor[(16), float32] */) /* ty=Tensor[(1, 16), float32] */;\n",
       "  %18 = nn.relu(%17) /* ty=Tensor[(1, 16), float32] span=aten__relu_3:0:0 */;\n",
       "  %19 = annotation.cast_hint(%18, dtype=\"int8\") /* ty=Tensor[(1, 16), float32] */;\n",
       "  %20 = annotation.stop_fusion(%19) /* ty=Tensor[(1, 16), float32] */;\n",
       "  %21 = nn.dense(%20, meta[relay.Constant][6] /* ty=Tensor[(8, 16), float32] */, units=None) /* ty=Tensor[(1, 8), float32] span=aten__linear_2:0:0 */;\n",
       "  add(%21, meta[relay.Constant][7] /* ty=Tensor[(8), float32] */) /* ty=Tensor[(1, 8), float32] */\n",
       "} /* ty=fn (Tensor[(1, 3, 4, 4), float32]) -> Tensor[(1, 8), float32] */"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "run_mod = relay.quantize.partition()(mod)\n",
    "run_mod[\"main\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## `dense+add+relu` 分区"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'tvm.relay.quantize._partition.QPartitionExpr'>\n",
      "<class 'tvm.relay.quantize._partition.QPartitionExpr'>\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "fn (%data: Tensor[(1, 3, 4, 4), float32] /* ty=Tensor[(1, 3, 4, 4), float32] span=aten___convolution_0_data:0:0 */) -> Tensor[(1, 8), float32] {\n",
       "  %0 = nn.conv2d(%data, meta[relay.Constant][0] /* ty=Tensor[(16, 3, 3, 3), float32] */, padding=[1, 1, 1, 1], channels=16, kernel_size=[3, 3]) /* ty=Tensor[(1, 16, 4, 4), float32] span=aten___convolution_0:0:0 */;\n",
       "  %1 = nn.relu(%0) /* ty=Tensor[(1, 16, 4, 4), float32] span=aten__relu_0:0:0 */;\n",
       "  %2 = annotation.cast_hint(%1, dtype=\"int8\") /* ty=Tensor[(1, 16, 4, 4), float32] */;\n",
       "  %3 = annotation.stop_fusion(%2) /* ty=Tensor[(1, 16, 4, 4), float32] */;\n",
       "  %4 = nn.conv2d(%3, meta[relay.Constant][1] /* ty=Tensor[(16, 16, 3, 3), float32] */, padding=[1, 1, 1, 1], channels=16, kernel_size=[3, 3]) /* ty=Tensor[(1, 16, 4, 4), float32] span=aten___convolution_1:0:0 */;\n",
       "  %5 = add(%4, meta[relay.Constant][2] /* ty=Tensor[(16, 1, 1), float32] */) /* ty=Tensor[(1, 16, 4, 4), float32] */;\n",
       "  %6 = nn.relu(%5) /* ty=Tensor[(1, 16, 4, 4), float32] span=aten__relu_1:0:0 */;\n",
       "  %7 = nn.max_pool2d(%6, pool_size=[2, 2], strides=[2, 2], padding=[0, 0, 0, 0]) /* ty=Tensor[(1, 16, 2, 2), float32] span=aten__max_pool2d_0:0:0 */;\n",
       "  %8 = annotation.cast_hint(%7, dtype=\"int8\") /* ty=Tensor[(1, 16, 2, 2), float32] */;\n",
       "  %9 = annotation.stop_fusion(%8) /* ty=Tensor[(1, 16, 2, 2), float32] */;\n",
       "  %10 = reshape(%9, newshape=[0, -1, 1, 1]) /* ty=Tensor[(1, 64, 1, 1), float32] span=aten__flatten_0:0:0 */;\n",
       "  %11 = squeeze(%10, axis=[2, 3]) /* ty=Tensor[(1, 64), float32] span=aten__flatten_0:0:0 */;\n",
       "  %12 = nn.dense(%11, meta[relay.Constant][3] /* ty=Tensor[(32, 64), float32] */, units=None) /* ty=Tensor[(1, 32), float32] span=aten__linear_0:0:0 */;\n",
       "  %13 = nn.relu(%12) /* ty=Tensor[(1, 32), float32] span=aten__relu_2:0:0 */;\n",
       "  %14 = annotation.cast_hint(%13, dtype=\"int8\") /* ty=Tensor[(1, 32), float32] */;\n",
       "  %15 = annotation.stop_fusion(%14) /* ty=Tensor[(1, 32), float32] */;\n",
       "  %16 = nn.dense(%15, meta[relay.Constant][4] /* ty=Tensor[(16, 32), float32] */, units=None) /* ty=Tensor[(1, 16), float32] span=aten__linear_1:0:0 */;\n",
       "  %17 = add(%16, meta[relay.Constant][5] /* ty=Tensor[(16), float32] */) /* ty=Tensor[(1, 16), float32] */;\n",
       "  %18 = nn.relu(%17) /* ty=Tensor[(1, 16), float32] span=aten__relu_3:0:0 */;\n",
       "  %19 = annotation.cast_hint(%18, dtype=\"int8\") /* ty=Tensor[(1, 16), float32] */;\n",
       "  %20 = annotation.stop_fusion(%19) /* ty=Tensor[(1, 16), float32] */;\n",
       "  %21 = nn.dense(%20, meta[relay.Constant][6] /* ty=Tensor[(8, 16), float32] */, units=None) /* ty=Tensor[(1, 8), float32] span=aten__linear_2:0:0 */;\n",
       "  %22 = add(%21, meta[relay.Constant][7] /* ty=Tensor[(8), float32] */) /* ty=Tensor[(1, 8), float32] */;\n",
       "  %23 = annotation.cast_hint(%22, dtype=\"int8\") /* ty=Tensor[(1, 8), float32] */;\n",
       "  annotation.stop_fusion(%23) /* ty=Tensor[(1, 8), float32] */\n",
       "} /* ty=fn (Tensor[(1, 3, 4, 4), float32]) -> Tensor[(1, 8), float32] */"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from tvm.relay import analysis as _analysis\n",
    "\n",
    "reset_partition()\n",
    "_op.get(\"nn.dense\").reset_attr(\"FQPartitionRewrite\")\n",
    "def add_partition_generic(ref_call, new_args, ctx):\n",
    "    \"\"\"Rewrite function for ewise add for partition for generic devices\"\"\"\n",
    "    lhs_cond, lhs = partition_expr_check(new_args[0])\n",
    "    rhs_cond, rhs = partition_expr_check(new_args[1])\n",
    "    if lhs_cond and rhs_cond:\n",
    "        # - introduced by ResNet, when for the first residual connection\n",
    "        #     ...\n",
    "        #     %0 = nn.conv2d(%data, %meta[relay.Constant])\n",
    "        #     %1 = add(%0, %meta[relay.Constant])\n",
    "        #     %2 = nn.relu(%1)\n",
    "        #     %3 = nn.max_pool2d(%2)\n",
    "        #     ...\n",
    "        #     %9 = nn.conv2d(%8, %meta[relay.Constant])\n",
    "        #     %10 = add(%9, %meta[relay.Constant])\n",
    "        #     %11 = add(%3, %10)  <- need to insert annotations for %3, %10\n",
    "        #     ...\n",
    "        lhs = new_args[0].realize()\n",
    "        rhs = new_args[1].realize()\n",
    "        return QPartitionExpr(_forward_op(ref_call, [lhs, rhs]))\n",
    "    if not lhs_cond and rhs_cond:\n",
    "        # - introduced by residual connection in ResNet\n",
    "        #     ...\n",
    "        #     %13 = nn.conv2d(%12, %meta[relay.Constant])\n",
    "        #     %14 = add(%13, %meta[relay.Constant])\n",
    "        #     %15 = annotation.cast_hint(%15, 'int8')\n",
    "        #     %16 = annotation.stop_fusion(%16)\n",
    "        #     %17 = add(%5, %16)\n",
    "        #     %18 = nn.relu(%17)\n",
    "        #     ...\n",
    "        #     %24 = nn.conv2d(%23, %meta[relay.Constant])\n",
    "        #     %25 = add(%24, %meta[relay.Constant])\n",
    "        #     %26 = add(%18, %25)  <- need to insert annotations for %25\n",
    "        #     ...\n",
    "        rhs = new_args[1].realize()\n",
    "        return _forward_op(ref_call, [lhs, rhs])\n",
    "    if lhs_cond and not rhs_cond:\n",
    "        if _analysis.check_constant(rhs):\n",
    "            # - introduced by batch_norm: add(out, bias)\n",
    "            return QPartitionExpr(_forward_op(ref_call, [lhs, rhs]))\n",
    "        # - introduced by residual connection in MobileNetV2\n",
    "        #     ...\n",
    "        #     %81 = add(%80, meta[relay.Constant])\n",
    "        #     %82 = annotation.cast_hint(%81, 'int8')\n",
    "        #     %83 = annotation.stop_fusion(%82)\n",
    "        #     %84 = add(%79, %83)\n",
    "        #     ...\n",
    "        #     %96 = nn.conv2d(%94, %meta[relay.Constant])\n",
    "        #     %96 = add(%95, %meta[relay.Constant])\n",
    "        #     %97 = add(%96, %84)  <- need to insert annotations for %96\n",
    "        #     ...\n",
    "        lhs = new_args[0].realize()\n",
    "        return _forward_op(ref_call, [lhs, rhs])\n",
    "    if not lhs_cond and not rhs_cond:\n",
    "        # trivial case\n",
    "        return None\n",
    "\n",
    "    raise ValueError\n",
    "@register_partition_function(\"nn.dense\")\n",
    "def dense_partition_function(ref_call, new_args, ctx):\n",
    "    \"\"\"Rewrite function for dense for partition\"\"\"\n",
    "    data_cond, data = partition_expr_check(new_args[0])\n",
    "    kernel_cond, kernel = partition_expr_check(new_args[1])\n",
    "    assert not kernel_cond\n",
    "    if data_cond:\n",
    "        print(type(new_args[0]))\n",
    "        data = new_args[0].realize()\n",
    "    ret = _forward_op(ref_call, [data, kernel])\n",
    "    return QPartitionExpr(ret)\n",
    "register_partition_function(\"nn.relu\", identity_partition_function)\n",
    "register_partition_function(\"nn.max_pool2d\", identity_partition_function)\n",
    "register_partition_function(\"add\", add_partition_generic)\n",
    "run_mod = relay.quantize.partition()(mod)\n",
    "run_mod[\"main\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'tvm.relay.quantize._partition.QPartitionExpr'>\n",
      "<class 'tvm.relay.quantize._partition.QPartitionExpr'>\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:autotvm:One or more operators have not been tuned. Please tune your model for better performance. Use DEBUG logging level to see more details.\n"
     ]
    }
   ],
   "source": [
    "def dataset():\n",
    "    for _ in range(1):\n",
    "        data = np.random.normal(size=(1, 3, 4, 4)).astype(\"float32\")\n",
    "        yield {\"data\": data}\n",
    "with tvm.transform.PassContext(opt_level=3):\n",
    "    with relay.quantize.qconfig(\n",
    "        calibrate_mode=\"kl_divergence\",\n",
    "        weight_scale=\"max\",\n",
    "        skip_conv_layers=[],\n",
    "        skip_dense_layer=False,\n",
    "    ):\n",
    "        qmod = relay.quantize.quantize(mod, params, dataset())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "def @main(%data: Tensor[(1, 3, 4, 4), float32] /* ty=Tensor[(1, 3, 4, 4), float32] span=aten___convolution_0_data:0:0 */) -> Tensor[(1, 8), float32] {\n",
      "  %44 = fn (%p08: Tensor[(1, 3, 4, 4), float32] /* ty=Tensor[(1, 3, 4, 4), float32] */, Primitive=1) -> Tensor[(1, 3, 4, 4), int8] {\n",
      "    %41 = multiply(%p08, 77.4988f /* ty=float32 */) /* ty=Tensor[(1, 3, 4, 4), float32] */;\n",
      "    %42 = round(%41) /* ty=Tensor[(1, 3, 4, 4), float32] */;\n",
      "    %43 = clip(%42, a_min=-127f, a_max=127f) /* ty=Tensor[(1, 3, 4, 4), float32] */;\n",
      "    cast(%43, dtype=\"int8\") /* ty=Tensor[(1, 3, 4, 4), int8] */\n",
      "  } /* ty=fn (Tensor[(1, 3, 4, 4), float32]) -> Tensor[(1, 3, 4, 4), int8] */;\n",
      "  %45 = %44(%data) /* ty=Tensor[(1, 3, 4, 4), int8] */;\n",
      "  %46 = fn (%p07: Tensor[(1, 3, 4, 4), int8] /* ty=Tensor[(1, 3, 4, 4), int8] */, %p14: Tensor[(16, 3, 3, 3), int8] /* ty=Tensor[(16, 3, 3, 3), int8] */, Primitive=1) -> Tensor[(1, 16, 4, 4), int8] {\n",
      "    %35 = nn.conv2d(%p07, %p14, padding=[1, 1, 1, 1], channels=16, kernel_size=[3, 3], out_dtype=\"int32\") /* ty=Tensor[(1, 16, 4, 4), int32] */;\n",
      "    %36 = nn.relu(%35) /* ty=Tensor[(1, 16, 4, 4), int32] */;\n",
      "    %37 = cast(%36, dtype=\"int64\") /* ty=Tensor[(1, 16, 4, 4), int64] */;\n",
      "    %38 = fixed_point_multiply(%37, multiplier=1263179136, shift=-8) /* ty=Tensor[(1, 16, 4, 4), int64] */;\n",
      "    %39 = clip(%38, a_min=-127f, a_max=127f) /* ty=Tensor[(1, 16, 4, 4), int64] */;\n",
      "    %40 = cast(%39, dtype=\"int32\") /* ty=Tensor[(1, 16, 4, 4), int32] */;\n",
      "    cast(%40, dtype=\"int8\") /* ty=Tensor[(1, 16, 4, 4), int8] */\n",
      "  } /* ty=fn (Tensor[(1, 3, 4, 4), int8], Tensor[(16, 3, 3, 3), int8]) -> Tensor[(1, 16, 4, 4), int8] */;\n",
      "  %47 = %46(%45, meta[relay.Constant][0] /* ty=Tensor[(16, 3, 3, 3), int8] */) /* ty=Tensor[(1, 16, 4, 4), int8] */;\n",
      "  %48 = fn (%p06: Tensor[(1, 16, 4, 4), int8] /* ty=Tensor[(1, 16, 4, 4), int8] */, %p13: Tensor[(16, 16, 3, 3), int8] /* ty=Tensor[(16, 16, 3, 3), int8] */, %p22: Tensor[(16, 1, 1), int32] /* ty=Tensor[(16, 1, 1), int32] */, Primitive=1) -> Tensor[(1, 16, 4, 4), int8] {\n",
      "    %28 = nn.conv2d(%p06, %p13, padding=[1, 1, 1, 1], channels=16, kernel_size=[3, 3], out_dtype=\"int32\") /* ty=Tensor[(1, 16, 4, 4), int32] */;\n",
      "    %29 = add(%28, %p22) /* ty=Tensor[(1, 16, 4, 4), int32] */;\n",
      "    %30 = nn.relu(%29) /* ty=Tensor[(1, 16, 4, 4), int32] */;\n",
      "    %31 = cast(%30, dtype=\"int64\") /* ty=Tensor[(1, 16, 4, 4), int64] */;\n",
      "    %32 = fixed_point_multiply(%31, multiplier=1588850048, shift=-9) /* ty=Tensor[(1, 16, 4, 4), int64] */;\n",
      "    %33 = clip(%32, a_min=-127f, a_max=127f) /* ty=Tensor[(1, 16, 4, 4), int64] */;\n",
      "    %34 = cast(%33, dtype=\"int32\") /* ty=Tensor[(1, 16, 4, 4), int32] */;\n",
      "    cast(%34, dtype=\"int8\") /* ty=Tensor[(1, 16, 4, 4), int8] */\n",
      "  } /* ty=fn (Tensor[(1, 16, 4, 4), int8], Tensor[(16, 16, 3, 3), int8], Tensor[(16, 1, 1), int32]) -> Tensor[(1, 16, 4, 4), int8] */;\n",
      "  %49 = %48(%47, meta[relay.Constant][1] /* ty=Tensor[(16, 16, 3, 3), int8] */, meta[relay.Constant][2] /* ty=Tensor[(16, 1, 1), int32] */) /* ty=Tensor[(1, 16, 4, 4), int8] */;\n",
      "  %50 = fn (%p05: Tensor[(1, 16, 4, 4), int8] /* ty=Tensor[(1, 16, 4, 4), int8] */, Primitive=1) -> Tensor[(1, 16, 2, 2), int8] {\n",
      "    %27 = nn.max_pool2d(%p05, pool_size=[2, 2], strides=[2, 2], padding=[0, 0, 0, 0]) /* ty=Tensor[(1, 16, 2, 2), int8] */;\n",
      "    cast(%27, dtype=\"int8\") /* ty=Tensor[(1, 16, 2, 2), int8] */\n",
      "  } /* ty=fn (Tensor[(1, 16, 4, 4), int8]) -> Tensor[(1, 16, 2, 2), int8] */;\n",
      "  %51 = %50(%49) /* ty=Tensor[(1, 16, 2, 2), int8] */;\n",
      "  %52 = fn (%p04: Tensor[(1, 16, 2, 2), int8] /* ty=Tensor[(1, 16, 2, 2), int8] */, Primitive=1) -> Tensor[(1, 64), int8] {\n",
      "    %20 = reshape(%p04, newshape=[0, -1, 1, 1]) /* ty=Tensor[(1, 64, 1, 1), int8] */;\n",
      "    %21 = cast(%20, dtype=\"float32\") /* ty=Tensor[(1, 64, 1, 1), float32] */;\n",
      "    %22 = multiply(%21, 0.00379081f /* ty=float32 */) /* ty=Tensor[(1, 64, 1, 1), float32] */;\n",
      "    %23 = squeeze(%22, axis=[2, 3]) /* ty=Tensor[(1, 64), float32] span=aten__flatten_0:0:0 */;\n",
      "    %24 = multiply(%23, 270.912f /* ty=float32 */) /* ty=Tensor[(1, 64), float32] */;\n",
      "    %25 = round(%24) /* ty=Tensor[(1, 64), float32] */;\n",
      "    %26 = clip(%25, a_min=-127f, a_max=127f) /* ty=Tensor[(1, 64), float32] */;\n",
      "    cast(%26, dtype=\"int8\") /* ty=Tensor[(1, 64), int8] */\n",
      "  } /* ty=fn (Tensor[(1, 16, 2, 2), int8]) -> Tensor[(1, 64), int8] */;\n",
      "  %53 = %52(%51) /* ty=Tensor[(1, 64), int8] */;\n",
      "  %54 = fn (%p03: Tensor[(1, 64), int8] /* ty=Tensor[(1, 64), int8] */, %p12: Tensor[(32, 64), int8] /* ty=Tensor[(32, 64), int8] */, Primitive=1) -> Tensor[(1, 32), int8] {\n",
      "    %14 = nn.dense(%p03, %p12, units=None, out_dtype=\"int32\") /* ty=Tensor[(1, 32), int32] */;\n",
      "    %15 = nn.relu(%14) /* ty=Tensor[(1, 32), int32] */;\n",
      "    %16 = cast(%15, dtype=\"int64\") /* ty=Tensor[(1, 32), int64] */;\n",
      "    %17 = fixed_point_multiply(%16, multiplier=1167015808, shift=-8) /* ty=Tensor[(1, 32), int64] */;\n",
      "    %18 = clip(%17, a_min=-127f, a_max=127f) /* ty=Tensor[(1, 32), int64] */;\n",
      "    %19 = cast(%18, dtype=\"int32\") /* ty=Tensor[(1, 32), int32] */;\n",
      "    cast(%19, dtype=\"int8\") /* ty=Tensor[(1, 32), int8] */\n",
      "  } /* ty=fn (Tensor[(1, 64), int8], Tensor[(32, 64), int8]) -> Tensor[(1, 32), int8] */;\n",
      "  %55 = %54(%53, meta[relay.Constant][3] /* ty=Tensor[(32, 64), int8] */) /* ty=Tensor[(1, 32), int8] */;\n",
      "  %56 = fn (%p02: Tensor[(1, 32), int8] /* ty=Tensor[(1, 32), int8] */, %p11: Tensor[(16, 32), int8] /* ty=Tensor[(16, 32), int8] */, %p21: Tensor[(16), int32] /* ty=Tensor[(16), int32] */, Primitive=1) -> Tensor[(1, 16), int8] {\n",
      "    %7 = nn.dense(%p02, %p11, units=None, out_dtype=\"int32\") /* ty=Tensor[(1, 16), int32] */;\n",
      "    %8 = add(%7, %p21) /* ty=Tensor[(1, 16), int32] */;\n",
      "    %9 = nn.relu(%8) /* ty=Tensor[(1, 16), int32] */;\n",
      "    %10 = cast(%9, dtype=\"int64\") /* ty=Tensor[(1, 16), int64] */;\n",
      "    %11 = fixed_point_multiply(%10, multiplier=1117116032, shift=-8) /* ty=Tensor[(1, 16), int64] */;\n",
      "    %12 = clip(%11, a_min=-127f, a_max=127f) /* ty=Tensor[(1, 16), int64] */;\n",
      "    %13 = cast(%12, dtype=\"int32\") /* ty=Tensor[(1, 16), int32] */;\n",
      "    cast(%13, dtype=\"int8\") /* ty=Tensor[(1, 16), int8] */\n",
      "  } /* ty=fn (Tensor[(1, 32), int8], Tensor[(16, 32), int8], Tensor[(16), int32]) -> Tensor[(1, 16), int8] */;\n",
      "  %57 = %56(%55, meta[relay.Constant][4] /* ty=Tensor[(16, 32), int8] */, meta[relay.Constant][5] /* ty=Tensor[(16), int32] */) /* ty=Tensor[(1, 16), int8] */;\n",
      "  %58 = fn (%p01: Tensor[(1, 16), int8] /* ty=Tensor[(1, 16), int8] */, %p1: Tensor[(8, 16), int8] /* ty=Tensor[(8, 16), int8] */, %p2: Tensor[(8), int32] /* ty=Tensor[(8), int32] */, Primitive=1) -> Tensor[(1, 8), int8] {\n",
      "    %1 = nn.dense(%p01, %p1, units=None, out_dtype=\"int32\") /* ty=Tensor[(1, 8), int32] */;\n",
      "    %2 = add(%1, %p2) /* ty=Tensor[(1, 8), int32] */;\n",
      "    %3 = cast(%2, dtype=\"int64\") /* ty=Tensor[(1, 8), int64] */;\n",
      "    %4 = fixed_point_multiply(%3, multiplier=1287755520, shift=-9) /* ty=Tensor[(1, 8), int64] */;\n",
      "    %5 = clip(%4, a_min=-127f, a_max=127f) /* ty=Tensor[(1, 8), int64] */;\n",
      "    %6 = cast(%5, dtype=\"int32\") /* ty=Tensor[(1, 8), int32] */;\n",
      "    cast(%6, dtype=\"int8\") /* ty=Tensor[(1, 8), int8] */\n",
      "  } /* ty=fn (Tensor[(1, 16), int8], Tensor[(8, 16), int8], Tensor[(8), int32]) -> Tensor[(1, 8), int8] */;\n",
      "  %59 = %58(%57, meta[relay.Constant][6] /* ty=Tensor[(8, 16), int8] */, meta[relay.Constant][7] /* ty=Tensor[(8), int32] */) /* ty=Tensor[(1, 8), int8] */;\n",
      "  %60 = fn (%p0: Tensor[(1, 8), int8] /* ty=Tensor[(1, 8), int8] */, Primitive=1) -> Tensor[(1, 8), float32] {\n",
      "    %0 = cast(%p0, dtype=\"float32\") /* ty=Tensor[(1, 8), float32] */;\n",
      "    multiply(%0, 0.00185884f /* ty=float32 */) /* ty=Tensor[(1, 8), float32] */\n",
      "  } /* ty=fn (Tensor[(1, 8), int8]) -> Tensor[(1, 8), float32] */;\n",
      "  %60(%59) /* ty=Tensor[(1, 8), float32] */\n",
      "}\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(relay.transform.FuseOps()(qmod))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "tvm-env",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.2"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
